//Udacity HW 4
//Radix Sorting

#include "utils.h"
#include <thrust/host_vector.h>

/* Red Eye Removal
   ===============
   
   For this assignment we are implementing red eye removal.  This is
   accomplished by first creating a score for every pixel that tells us how
   likely it is to be a red eye pixel.  We have already done this for you - you
   are receiving the scores and need to sort them in ascending order so that we
   know which pixels to alter to remove the red eye.

   Note: ascending order == smallest to largest

   Each score is associated with a position, when you sort the scores, you must
   also move the positions accordingly.

   Implementing Parallel Radix Sort with CUDA
   ==========================================

   The basic idea is to construct a histogram on each pass of how many of each
   "digit" there are.   Then we scan this histogram so that we know where to put
   the output of each digit.  For example, the first 1 must come after all the
   0s so we have to know how many 0s there are to be able to start moving 1s
   into the correct position.

   1) Histogram of the number of occurrences of each digit
   2) Exclusive Prefix Sum of Histogram
   3) Determine relative offset of each digit
        For example [0 0 1 1 0 0 1]
                ->  [0 1 0 1 2 3 2]
   4) Combine the results of steps 2 & 3 to determine the final
      output location for each element and move it there

   LSB Radix sort is an out-of-place sort and you will need to ping-pong values
   between the input and output buffers we have provided.  Make sure the final
   sorted results end up in the output buffer!  Hint: You may need to do a copy
   at the end.
 */

const int BLOCK_SIZE = 256; 
const int UINT32_BITS = 32;  // 32位无符号整数的位数

// 扫描类型枚举
enum ScanType {
 INCLUSIVE,  // 包含式扫描(包含当前元素)
 EXCLUSIVE   // 排除式扫描(不包含当前元素)
};

__global__
void naive_histogram(unsigned int* const d_histo, int numBins,
                     unsigned int* const d_in, int size, int shift) {
  int i = threadIdx.x + blockDim.x * blockIdx.x;
  if (i >= size) return;

  unsigned int bit = ((d_in[i] >> shift) & 1u);
  atomicAdd(&d_histo[bit], 1);
}

// Hillis-Steele并行扫描核函数
__global__
void scan_hillis_steele(unsigned int* d_out, const unsigned int* d_in, int size, ScanType scan_type) {
  extern __shared__ unsigned int shared_data[];  // 两倍 blocksize 大小,双缓冲避免读写冲突
  int tid = threadIdx.x;  // 线程ID
  int globalId = blockDim.x * blockIdx.x + threadIdx.x;  // 全局ID
  if (globalId >= size) return;  // 边界检查

  // 根据扫描类型初始化
  if (scan_type == INCLUSIVE) {
    shared_data[tid] = d_in[globalId];  // 包含式
  } else {
    shared_data[tid] = (globalId == 0 ? 0 : d_in[globalId-1]);  // 排除式
  }
  __syncthreads();  // 同步线程

  // Hillis-Steele扫描算法
  int pin = 0, pout = 1;  // 双缓冲索引
  for (int s = 1; s < blockDim.x; s <<= 1) {  // 步长倍增
    if (tid < s) {
      shared_data[pout * blockDim.x + tid] = shared_data[pin * blockDim.x + tid];
    }
    if (tid + s < blockDim.x) {
      shared_data[pout * blockDim.x + tid + s] = 
        shared_data[pin * blockDim.x + tid] + shared_data[pin * blockDim.x + tid + s];
    }
    __syncthreads();
    pin = 1 - pin; pout = 1 - pout;  // 交换缓冲区
  }
  d_out[globalId] = shared_data[pin * blockDim.x + tid];  // 写入结果
}

// 块扫描相加核函数
__global__
void block_scan_add_block_sum(unsigned int* d_out, const unsigned int* const block_sum, 
                             const unsigned int* const d_scan, int size) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;  // 全局索引
  if (i >= size) return;  // 边界检查
  d_out[i] = block_sum[blockIdx.x] + d_scan[i];  // 合并结果
}

// 零位映射核函数
__global__
void zero_digit_mapper(unsigned int* d_out, unsigned int* const d_inputVals,
                      size_t numElems, int shift) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;  // 全局索引
  if (i >= numElems) return;  // 边界检查
  d_out[i] = !((d_inputVals[i] >> shift) & 1u);  // 检查是否为0位,为0则标记为1
}

// 初始化块和核函数
__global__ 
void init_block_sum(unsigned int* block_sum, const unsigned int* const block_level_scan, int block_num) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;  // 全局索引
  if (i < block_num) {  // 边界检查
    block_sum[i] = (i == 0) ? 0 : block_level_scan[i * blockDim.x - 1];  // 计算块前缀和
  }
}

// 递归扫描函数(如果递归次数过多，存在内存问题)
void scan(unsigned int* d_out, const unsigned int* d_in, int size, ScanType scan_type, int recursive){
  if (size <= 0) return;  // 边界检查
  // printf("recursive = %d\n", recursive);

  // 1. 块级扫描
  int block_num = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;  // 计算块数
  scan_hillis_steele<<<block_num, BLOCK_SIZE, 2 * BLOCK_SIZE * sizeof(unsigned int)>>>(d_out, d_in, size, scan_type);
  cudaDeviceSynchronize();

  if (block_num == 1) {
    return;  // 单块直接返回
  }

  // 2. 计算块和并递归扫描
  unsigned int* d_block_sum;
  cudaMalloc(&d_block_sum, sizeof(unsigned int) * block_num);
  int block_num2 = (block_num + BLOCK_SIZE-1)/BLOCK_SIZE;
  init_block_sum<<<block_num2, BLOCK_SIZE>>>(d_block_sum, d_out, block_num);

  scan(d_block_sum, d_block_sum, block_num, INCLUSIVE, recursive + 1);

   // 3. 合并扫描结果
  block_scan_add_block_sum<<<block_num, BLOCK_SIZE>>>(d_out, d_block_sum, d_out, size);
  cudaDeviceSynchronize(); 

  cudaFree(d_block_sum);  // 释放内存
}

// 数据移动核函数
__global__ 
void move_data(unsigned int* d_out_vals, unsigned int* d_out_pos,
              unsigned int* const d_in_values, unsigned int* const d_in_pos,
              unsigned int* d_zero_pos, unsigned int* d_histo,
              int shift, int size) {
  int i = threadIdx.x + blockDim.x * blockIdx.x;  // 全局索引
  if (i >= size) return;  // 边界检查

  unsigned int tmp_value = d_in_values[i];  // 临时存储值
  unsigned int tmp_pos = d_in_pos[i];  // 临时存储位置
  __syncthreads();  // 同步线程

  unsigned int num_zeros = d_histo[0];  // 获取0位总数量

  // 根据当前位值移动数据
  if(((d_in_values[i] >> shift) & 0x1u) == 0) {
    d_out_vals[d_zero_pos[i]] = tmp_value;  // 0位放置在前部
    d_out_pos[d_zero_pos[i]] = tmp_pos;
  } else {
    // 1位放置在所有0位的后部
    d_out_vals[i - d_zero_pos[i] + num_zeros] = tmp_value;
    d_out_pos[i - d_zero_pos[i] + num_zeros] = tmp_pos;
  }
}

void gpu_radix_sort(unsigned int* const d_inputVals,
               unsigned int* const d_inputPos,
               unsigned int* const d_outputVals,
               unsigned int* const d_outputPos,
               const size_t numElems){
  // 分配设备内存用于存储0位标记
  unsigned int* d_zero;
  cudaMalloc(&d_zero, sizeof(unsigned int)*numElems);

  // 分配设备内存用于直方图(2 bins)
  unsigned int* d_histo;
  cudaMalloc(&d_histo, sizeof(unsigned int)*2);

  // 初始化源指针
  unsigned int* d_vals_src = d_inputVals;
  unsigned int* d_pos_src = d_inputPos;

  // 初始化目标指针
  unsigned int* d_vals_dst_tmp = d_outputVals;
  unsigned int* d_pos_dst_tmp = d_outputPos;

  int gridsize = (numElems + BLOCK_SIZE - 1)/BLOCK_SIZE;

  // 对32位无符号整数的每一位进行处理(从最低位开始)
  for (int shift=0; shift < UINT32_BITS; shift++) {
    // 重置直方图
    cudaMemset(d_histo, 0, sizeof(unsigned int)*2);
    
    // 计算当前位的直方图(0和1的计数)
    naive_histogram<<<gridsize, BLOCK_SIZE>>>(d_histo, 2, d_vals_src, numElems, shift);
    cudaDeviceSynchronize();

    // 标记当前位为0的元素,为0则标记为1
    zero_digit_mapper<<<gridsize, BLOCK_SIZE>>>(d_zero, d_vals_src, numElems, shift);
    cudaDeviceSynchronize(); 
 
    // 对0位标记进行前缀和扫描(非包含)
    scan(d_zero, d_zero, numElems, EXCLUSIVE, 1);
  
    // 根据扫描结果移动数据到正确位置
    move_data<<<gridsize, BLOCK_SIZE>>>(d_vals_dst_tmp, d_pos_dst_tmp,
                                                                  d_vals_src, d_pos_src, d_zero, d_histo, shift, numElems);
    cudaDeviceSynchronize(); 

    // 缓冲区交换
    std::swap(d_vals_src, d_vals_dst_tmp);
    std::swap(d_pos_src, d_pos_dst_tmp);
  }

  // 因为每次 swap 后，src 总是指向最新结果
  if(UINT32_BITS % 2 == 0){
    cudaMemcpy(d_outputVals, d_vals_src,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToDevice);
    cudaMemcpy(d_outputPos, d_pos_src,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToDevice);
  }

  // 释放临时内存
  cudaFree(d_histo);
  cudaFree(d_zero);
}

 void your_sort(unsigned int* const d_inputVals,
               unsigned int* const d_inputPos,
               unsigned int* const d_outputVals,
               unsigned int* const d_outputPos,
               const size_t numElems)
{ 
  //TODO PUT YOUR SORT HERE
  gpu_radix_sort(d_inputVals, d_inputPos, d_outputVals, d_outputPos, numElems);
}

#include <cub/cub.cuh>

// 使用cub库的radix排序
void cub_your_sort(unsigned int* const d_inputVals,
               unsigned int* const d_inputPos,
               unsigned int* const d_outputVals,
               unsigned int* const d_outputPos,
               const size_t numElems)
{ 
  void* d_temp_storage = nullptr;
  size_t temp_storage_bytes = 0;

  // 查询临时存储大小
  cub::DeviceRadixSort::SortPairs(
      d_temp_storage, temp_storage_bytes,
      d_inputVals, d_outputVals,
      d_inputPos, d_outputPos,
      numElems
  );

  // 分配临时内存
  cudaMalloc(&d_temp_storage, temp_storage_bytes);

  cub::DeviceRadixSort::SortPairs(
      d_temp_storage, temp_storage_bytes,
      d_inputVals, d_outputVals,
      d_inputPos, d_outputPos,
      numElems
  );

  cudaDeviceSynchronize();
  cudaFree(d_temp_storage);
}
